import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory

from Embed_model import BGE_Embed
from qw_model import qw_model
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

# 初始化模型和嵌入
model = qw_model()
embedding = BGE_Embed()

# 加载知识库
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_db'
)

# 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={'k': 3})


contextualize_q_system_prompt = (
    "给定聊天记录和最新用户问题 "
    "这可能会引用聊天历史中的上下文， "
    "制定一个可以理解的独立问题 "
    "没有聊天记录。不要回答问题， "
    "如果需要，只需重新规划，否则按原样退回。"
)
contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)
history_aware_retriever = create_history_aware_retriever(
    model, retriever, contextualize_q_prompt
)

system_prompt = (
    " 你是一个聪明又有亲和力的销售客服，名叫银弓，擅长用轻松、亲切的语气解答用户关于公司业务、价格、产品咨询的问题。你的目标是提供专业、准确的回答，尤其是价格问题，严格按照相应文档回答，风格需要口语化、自然流畅、贴近用户生活，避免书面化和过于正式的表达以及尊称。"
    "1. 当用户问题不清楚时，主动询问用户具体需求。"
    "2. 回答价格问题，严格按照检索的文档给出，不可自行编造。"
    "3. 当没有用户所说的地区，就按它的省份价格算，当省份也没有那就只能做不限地区的"
    "4. 保持沟通语气亲切、自然，避免过于正式的措辞，多用口语化的表达和轻松的语气，比如“老板”、“宝子”、“咱们”，禁止使用'您'。"
    "5. 输出答案应简洁明了，字数控制在15字左右。"
    "\n\n"
    "{context}"
)

# 创建提示词模板
qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

question_answer_chain = create_stuff_documents_chain(model, qa_prompt)

rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

# 有状态地管理聊天记录
store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

Final_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer",
)

# 创建 FastAPI 应用
app = FastAPI()

class ChatRequest(BaseModel):
    input: str
    session_id: str

@app.post("/chat/")
async def chat(request: ChatRequest):
    response = Final_chain.invoke(
        {"input": request.input},
        config={"configurable": {"session_id": request.session_id}}
    )["answer"]
    return {"answer": response}

# 运行应用
# uvicorn app:app --reload
if __name__ == '__main__':
    uvicorn.run(app, host="0.0.0.0", port=8000)

#curl -X 'POST' \ 'http://127.0.0.1:8000/chat/' \ -H
# 'Content-Type: application/json' \ -d '{ "input": "你好", "session_id": "12345" }'